# -*- coding: utf-8 -*-
"""
Created on Sun Oct 13 13:21:00 2013

@author: jkwong
"""
from sys import exit
import matplotlib.pyplot as plt
import numpy as np
# choose a combo!
comboMatch = [1, 2, 3]

# BUILD THE W MATRIX AND DISPLAY RESULTS

wMatrix = np.zeros((len(wMeanAllAll), len(comboMatch)))
countRangeMean = np.zeros(len(wMeanAllAll))
print statsCollapseListALL
print('Combo, count index, success, specificity, sensitivity')
for (countIndex, countRange) in enumerate(countRangesList):
    for (comboIndex, combo) in enumerate(comboList):
        if combo == comboMatch:
            wMatrix[countIndex,:] = wMeanAllAll[countIndex][comboIndex]
            print('%s, %d, %3.3f, %3.3f, %3.3f' %(str(combo), countIndex, successRateMeanAll[countIndex][comboIndex], \
                specificityMeanAll[countIndex][comboIndex], sensitivityMeanAll[countIndex][comboIndex]))
    countRangeMean[countIndex] = mean(countRange)


# FIT GAUSSIANS TO THE PB-ALL BAND
countBinEdges = np.arange(0, 375, 25)
countBinCenters = []
countBinEdgesList = []
for i in xrange(len(countBinEdges)-1):
    countBinEdgesList.append(np.array((countBinEdges[i], countBinEdges[i+1])))
    countBinCenters.append( (countBinEdges[i] + countBinEdges[i+1])/2.0)
countBinCenters = np.array(countBinCenters)

fitParamsPbAll = np.zeros( (len(countBinEdgesList), 3))
fitParamsPbNot = np.zeros( (len(countBinEdgesList), 3))

# choose the best w to use
wBEST = wMatrix[3,:]


# Calculate the distance quantity

histCountsPbAllList = []
histCountsPbNotList = []

dist0Best = np.sum(np.matlib.repmat(wBEST, statsMatrix.shape[0], 1) * \
            statsMatrix[:,comboMatch], axis = 1)
distBinedges = np.arange(-150, 50, 1)
distBincenters = (distBinedges[0:-1] + distBinedges[1:])/2.0

for (countBinEdgesIndex, countBinEdges) in enumerate(countBinEdgesList):
    #  PbAll
    # cut in counts
    cutt = (countMatrix >= countBinEdges[0]) & (countMatrix < countBinEdges[1])
    cutt = cutt & statsMatrix[:,-1].astype(bool)
    # make histogram of the values
    counts, binedges = np.histogram(dist0Best[cutt], distBinedges)
    histCountsPbAllList.append(counts)
    
    startingParam = [0.5 * max(counts), np.mean(dist0Best[cutt]), np.std(dist0Best[cutt])]
    try:
        popt, pcov = curve_fit(gauss_function, distBincenters, counts, p0 = startingParam)
    except:
        print("Gaussian, Bad Fit")
        popt = [-1, -1, -1]
    fitParamsPbAll[countBinEdgesIndex,:] = np.array(popt)
    
    #  PbNot
    # cut in counts
    cutt = (countMatrix >= countBinEdges[0]) & (countMatrix < countBinEdges[1])
    cutt = cutt & ~statsMatrix[:,-1].astype(bool)
    # make histogram of the values
    counts, binedges = np.histogram(dist0Best[cutt], distBinedges)
    histCountsPbNotList.append(counts)
    
    startingParam = [0.5 * max(counts), np.mean(dist0Best[cutt]), np.std(dist0Best[cutt])]
    try:
        popt, pcov = curve_fit(gauss_function, distBincenters, counts, p0 = startingParam)
    except:
        print("Gaussian, Bad Fit")
        popt = [-1, -1, -1]
    fitParamsPbNot[countBinEdgesIndex,:] = np.array(popt)

# FIT POLYNOMIAL
pfitmeanPbAll = np.polyfit(countBinCenters, fitParamsPbAll[:,1], 3)
pfitsigmaPbAll = np.polyfit(countBinCenters, fitParamsPbAll[:,2], 3)

pfitmeanPbNot = np.polyfit(countBinCenters, fitParamsPbAll[:,1], 3)
pfitsigmaPbNot = np.polyfit(countBinCenters, fitParamsPbAll[:,2], 3)



#
## PLOT W values as a function of count
#plt.figure()
#plt.grid()
#plt.plot(countRangeMean, wMatrix)
#plt.legend(statsCollapseListALL[comboMatch])
#plt.xlabel('Mean Count')
#plt.ylabel('Coefficient Value')
#plt.show()
#
## PLOT W values times the mean of the features as a function of count
#plt.figure()
#plt.grid()
#plt.plot(countRangeMean, wMatrix * np.tile(statsMatrix[:,comboMatch].mean(0), (wMatrix.shape[0], 1)))
#plt.legend(statsCollapseListALL[comboMatch])
#plt.xlabel('Mean Count')
#plt.ylabel('Coefficient Value * Mean Value')
#plt.show()
#
# FIT MEAN VS COUNTS

plt.figure()
plt.grid()
plt.plot(countBinCenters, fitParamsPbAll[:,1], 'xb', )
plt.plot(countBinCenters, fitParamsPbNot[:,1], 'xr')

x = np.arange(0, 400)
plt.plot(x, np.polyval(pfitmeanPbAll, x), '-k')
plt.xlabel('Count')
plt.ylabel('Dist Mean')
#
#
## FIT SIGMA VS COUNTS
#plt.figure()
#plt.grid()
#plt.plot(countBinCenters, fitParamsPbAll[:,2], 'xk')
#x = np.arange(0, 400)
#plt.plot(x, np.polyval(pfitsigmaPbAll, x), '-k')
#plt.xlabel('Count')
#plt.ylabel('Dist Sigma')

# plot histograms

index = 5

for index in xrange(len(histCountsPbAllList)):
    
    plt.figure()
    plt.grid()
    
    plt.bar(distBinedges[:-1], histCountsPbAllList[index], width = 1, color = 'b', edgecolor = 'b', fill = True, alpha = 0.6, label = 'PbALL')
    plt.bar(distBinedges[:-1], histCountsPbNotList[index], width = 1, color = 'r', edgecolor = 'r', fill = True, alpha = 0.6, label = 'PBNOT')
    
    plt.xlabel('Discriminant')
    plt.ylabel("Count")
    plt.title('Count Range: [%d, %d)' %(countBinEdgesList[index][0], countBinEdgesList[index][1]))



# SCATTER PLOT OF DIST VS COUNT AND THE BEST FIT LINE using best w for data points of all counts

plt.figure()
plt.grid()
for i in xrange(2):
    cutt = statsMatrix[:,-1] == i
    plt.plot(countMatrix[cutt], dist0Best[cutt], '.k', color = plotColors[i], alpha = 0.4, markersize = 10, label = groupNamesExportList[i])
plt.plot(countBinCenters, fitParamsPbAll[:,1], '-xg', markersize = 10)
plt.plot(countBinCenters, fitParamsPbAll[:,1] + 2.5*fitParamsPbAll[:,2], '-xg', markersize = 10, linewidth = 2)
plt.plot(countBinCenters, fitParamsPbAll[:,1] - 2.5*fitParamsPbAll[:,2], '-xg', markersize = 10, linewidth = 2)
x = np.arange(0, 400)
plt.plot(x, np.polyval(pfitmeanPbAll, x), '--g', linewidth = 3)
plt.legend()
plt.xlabel('Count')
plt.ylabel('Distance')
plt.title('')


# SCATTER PLOT OF DIST VS COUNT AND THE BEST FIT LINE, OFFSET Removed
plt.figure()
plt.grid()
for i in xrange(2):
    cutt = statsMatrix[:,-1] == i
    plt.plot(countMatrix[cutt], dist0Best[cutt] - np.polyval(pfitmeanPbAll, countMatrix[cutt]), '.k', \
    color = plotColors[i], alpha = 0.4, markersize = 10, label = groupNamesExportList[i])
##plt.plot(countBinCenters, fitParamsPbAll[:,1], '-xg', markersize = 10)
plt.plot(countBinCenters,+2.5*fitParamsPbAll[:,2], '-xg', markersize = 10, linewidth = 2)
plt.plot(countBinCenters, -2.5*fitParamsPbAll[:,2], '-xg', markersize = 10, linewidth = 2)
plt.legend()
plt.xlabel('Count')
plt.ylabel('Distance - Offset')
plt.title('')


# Plot of the Discriminant Versus Count, each figure corresponds to training at a particularly count range.
# Including Low Stats
for (countIndex, countRange) in enumerate(countRangesList):
    for (comboIndex, combo) in enumerate(comboList):
        if combo == comboMatch:
            print(statsCollapseListALL[combo])
            wMean = wMeanAllAll[countIndex][comboIndex]
            dist0 = np.sum(np.matlib.repmat(wMean, statsMatrixIncludingLowStats.shape[0], 1) * \
                        statsMatrixIncludingLowStats[:,combo], axis = 1)
            plt.figure()
            plt.grid()
            for i in xrange(2):
                cutt = statsMatrixIncludingLowStats[:,-1] == i
                cutt = cutt & (countMatrixIncludingLowStats[:,0] > 0)
                plt.plot(countMatrixIncludingLowStats[cutt,0], dist0[cutt], '.k', \
                color = plotColors[i], alpha = 0.4, markersize = 10, label = groupNamesExportList[i])
            plt.xlabel('Count')
            plt.ylabel('Distance')
            plt.title('Count Range: %s' % str(countRange))
            plt.legend()
            # draw bounds in count
            ylimit = plt.ylim()
            plt.plot([countRange[0], countRange[0]], [ylimit[0], ylimit[1]], '-g', linewidth = 1.5)
            plt.plot([countRange[1], countRange[1]], [ylimit[0], ylimit[1]], '-g', linewidth = 1.5)
            plt.ylim(ylimit)
plt.show()
exit()